#! /usr/bin/env python3

import pyspark.ml.feature
from pyspark.sql import SparkSession

from gai.v2.unify.dtype import SparkDataFrame
from gai.v2.unify.transformer.variant import SparkTransformerWrapper


class SQLTransformer(SparkTransformerWrapper):
    """Wrapper of Spark's ``SQLTransformer``.

    See Also:
        `pyspark.ml.feature.SQLTransformer <https://spark.apache.org/docs/latest/api/python/pyspark.ml.html#pyspark.ml.feature.SQLTransformer>`_

    Args:
        statement: the SQL statement
        instance_name: the name of this particular instance of ``Transformer``. It is recommended
            to give a unique name.
        pre_hooks: the list of hook functions to be called at the entry
            of method ``transform()``
        post_hooks: the list of hook functions to be called at the exit
            of method ``transform()``
        type_adaptive: indicator of whether incompatible dataset is
            automatically adapted


    >>> from gai.v2.unify import SparkDataFrame
    >>> from gai.v2.unify.transformer.sql_transformer import SQLTransformer
    >>> from pyspark.sql import SparkSession
    >>> spark = SparkSession.builder.master("local").getOrCreate()
    >>> df = spark.createDataFrame([("Alice", 12), ("Bob", 13), ("John", 23)],
    ...                            schema=("name", "age"))
    >>> ds = SparkDataFrame(df)
    >>> transformer = SQLTransformer("select * from __THIS__ a where a.age < 18", "find-less-than-18")
    >>> ds_res = transformer.transform(ds)
    >>> ds_res.unwrap().show()
    +-----+---+
    | name|age|
    +-----+---+
    |Alice| 12|
    |  Bob| 13|
    +-----+---+
    <BLANKLINE>
    >>> transformer.getExtraInfo()
    'select * from __THIS__ a where a.age < 18'
    >>> transformer.getKernelModule()
    'pyspark.ml.feature'
    """

    def __init__(self, statement=None, instance_name=None, pre_hooks=(), post_hooks=(), type_adaptive=False):
        transformer = pyspark.ml.feature.SQLTransformer(statement=statement)
        super(SQLTransformer, self).__init__(transformer=transformer,
                                             instance_name=instance_name,
                                             pre_hooks=pre_hooks,
                                             post_hooks=post_hooks,
                                             type_adaptive=type_adaptive)

    def getExtraInfo(self):
        """

        Returns:
            the SQL statement passed into this instance

        Notes:
            The return value of this method will appear as a field in the log.
        """
        return super().unwrap().getOrDefault("statement")
